#!/usr/bin/python3
# -*- coding: utf-8 -*-

"""
# @version: v1.0
# @author : wlisingle
# @Email : 19688513@qq.com
# @Project : trade-smart
# @File : stock_trading_strategy.py
# @Software: PyCharm
# @time: 2025/4/18 17:05
# @description : 股票交易策略模块
# 实现基于量价均线的交易信号生成和可视化
"""

import datetime
import pandas as pd
import matplotlib
import matplotlib.pyplot as plt

from matplotlib.dates import DateFormatter
from history_data.stock_history_manager import StockHistoryManager
from typing import Dict


class StockTradingStrategy:
    def __init__(self, data_manager):
        """初始化交易策略"""
        self.data_manager = data_manager

    def generate_signals(self, df: pd.DataFrame) -> pd.DataFrame:
        """生成买入和卖出信号"""
        signals = df.copy()

        signals['price'] = df['close']
        signals['volume'] = (df['volume'] / 10000).round(2)

        # 计算 RSI
        delta = df['close'].diff()
        gain = (delta.where(delta > 0, 0)).rolling(window=14).mean()
        loss = (-delta.where(delta < 0, 0)).rolling(window=14).mean()
        rs = gain / loss
        signals['rsi'] = 100 - (100 / (1 + rs))

        # 打印 RSI 值以进行调试
        print("RSI Values:", signals['rsi'].tail(10))

        # 生成买入信号条件

        # 信号1：成交量双均线条件
        signals['signal_vol'] = (
                (abs(signals['volume'] - signals['vol_ma5']) / signals['vol_ma5'].clip(lower=1e-6) < 0.02)
                &
                (signals['volume'] < signals['vol_ma60'] * 0.7)
        ).fillna(False)

        # 信号2：价格均线多头排列（增加斜率验证）
        signals['ma5_slope'] = df['ma5'].diff(3) > 0  # 3日斜率
        signals['signal_price'] = (
                (df['ma5'] > df['ma10']) &
                (df['ma10'] > df['ma20']) &
                (df['close'] > df['ma20']) &
                signals['ma5_slope']
        ).fillna(False)

        # 信号3：量能趋势突破（修正边界条件）
        signals['signal_vol_break'] = (
                (df['vol_ma5'] > df['vol_ma20']) &
                (df['vol_ma5'].shift(1) <= df['vol_ma20'].shift(1)) &
                (df['vol_ma5'] > df['vol_ma5'].shift(3))  # 确认趋势强度
        ).fillna(False)

        # 信号4：趋势延续验证（原signal_trend优化）
        signals['signal_trend'] = (
                (df['vol_ma5'] > df['vol_ma20']) &
                (df['vol_ma20'] > df['vol_ma60']) &
                (df['close'] > df['close'].rolling(3).mean())  # 3日动态支撑
        ).fillna(False)

        # 信号5：量价齐升（增强版）
        signals['signal_volume_price'] = (
                (df['volume'] > df['vol_ma5'] * 1.1) &  # 成交量显著放大
                (df['close'] > df['open']) &
                (df['close'] > df['high'].shift(1))  # 突破前日高点
        ).fillna(False)

        # 信号6：波动率过滤（滚动窗口计算）
        atr = (df['high'] - df['low']).rolling(14).mean()
        q30 = atr.rolling(252).quantile(0.3)  # 一年窗口
        q70 = atr.rolling(252).quantile(0.7)
        signals['signal_volatility'] = atr.between(q30, q70).astype(bool)

        # 信号7：RSI动态超卖（结合布林带）
        rsi_ma = signals['rsi'].rolling(14).mean()
        rsi_std = signals['rsi'].rolling(14).std()
        signals['signal_rsi'] = (
                signals['rsi'] < (rsi_ma - 1.5 * rsi_std)
        ).fillna(False)

        # 输出信号触发日期（修正列名）
        print("== 信号触发统计 ==")
        for signal_col in ['signal_vol', 'signal_price', 'signal_vol_break',
                           'signal_trend', 'signal_volume_price', 'signal_volatility', 'signal_rsi']:
            trigger_dates = signals[signals[signal_col]].index.strftime('%Y-%m-%d').tolist()
            print(f"{signal_col}: {len(trigger_dates)}日触发 → 示例日期: {trigger_dates[:3]}")

        return signals

    def export_signals_to_excel(self, all_signals: Dict[str, pd.DataFrame], BATCH_SIZE=500, matype=''):
        """高效分批次导出Excel文件（顺序修复版+格式优化）"""
        from tqdm import tqdm
        import math
        import os
        import datetime

        # 配置参数
        REQUIRED_COLUMNS = [
            'symbol', 'stock_name', 'strategy_date', 'price', 'volume', 'vol_ma5',
            'vol_ma10', 'vol_ma20', 'vol_ma30', 'vol_ma60', 'vol_ma120',
            'signal_vol', 'signal_price', 'signal_vol_break',
                           'signal_trend', 'signal_volume_price', 'signal_volatility', 'signal_rsi'
        ]
        CONDITION_COLUMNS = ['signal_vol', 'signal_price', 'signal_vol_break',
                           'signal_trend', 'signal_volume_price', 'signal_volatility', 'signal_rsi']  # 新增条件列定义

        # 预处理数据（保持原始顺序）
        signal_items = []
        for code in tqdm(sorted(all_signals.keys()), desc="预处理数据"):
            df = all_signals[code]
            if df.empty:
                continue

            # 列过滤和复制
            valid_cols = [col for col in REQUIRED_COLUMNS if col in df.columns]
            filtered_df = df[valid_cols].copy()

            # 日期格式预处理（修复字段名不一致问题）
            if 'strategy_date' in filtered_df.columns:
                filtered_df['strategy_date'] = pd.to_datetime(filtered_df['strategy_date']).dt.strftime('%Y-%m-%d')

            # 强制转换条件列为布尔类型（关键修复）
            for col in CONDITION_COLUMNS:
                if col in filtered_df.columns:
                    filtered_df[col] = filtered_df[col].astype(bool)

            signal_items.append((code, filtered_df))

        # 分批次处理
        total_batches = math.ceil(len(signal_items) / BATCH_SIZE)

        for batch_idx in tqdm(range(total_batches), desc="导出进度"):
            start = batch_idx * BATCH_SIZE
            end = min(start + BATCH_SIZE, len(signal_items))
            batch = signal_items[start:end]

            timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
            file_name = f"股票{matype}交易信号_{timestamp}_批次{batch_idx + 1:03d}.xlsx"

            try:
                with pd.ExcelWriter(
                        file_name,
                        engine='xlsxwriter',
                        engine_kwargs={'options': {'strings_to_numbers': True}}
                ) as writer:
                    workbook = writer.book

                    # 预定义格式（提升性能）
                    header_format = workbook.add_format({
                        'bold': True,
                        'bg_color': '#FFFF00',
                        'align': 'center',
                        'border': 1
                    })
                    text_format = workbook.add_format({'num_format': '@'})
                    condition_format = workbook.add_format({
                        'bg_color': '#C6EFCE',
                        'font_color': '#006100'
                    })

                    for idx, (stock_code, signals) in enumerate(batch):
                        sheet_name = f"{stock_code}_{idx}"[:31].translate(
                            str.maketrans('\\/*?:[]', '_______')
                        )

                        # 写入数据（表头与数据分离）
                        signals.to_excel(
                            writer,
                            sheet_name=sheet_name,
                            index=False,
                            header=False,
                            startrow=1
                        )

                        worksheet = writer.sheets[sheet_name]

                        # 写入自定义格式表头
                        for col_num, value in enumerate(signals.columns):
                            worksheet.write(0, col_num, value, header_format)

                        # 设置基础列格式
                        worksheet.set_column('A:Z', 15, text_format)

                        # ============== 条件格式修正版 ==============
                        for col_name in CONDITION_COLUMNS:
                            if col_name not in signals.columns:
                                continue

                            col_pos = signals.columns.get_loc(col_name)
                            start_row = 1  # 数据起始行（Excel第2行）
                            end_row = start_row + len(signals)  # 实际结束行号

                            # 精准计算范围（注意Excel行号从0开始）
                            worksheet.conditional_format(
                                first_row=start_row,
                                first_col=col_pos,
                                last_row=end_row - 1,  # 修正行号越界问题
                                last_col=col_pos,
                                options={
                                    'type': 'cell',
                                    'criteria': '==',
                                    'value': True,
                                    'format': condition_format
                                }
                            )
                        # ============================================

                        # 冻结首行（必须在最后设置）
                        worksheet.freeze_panes(1, 0)

                    print(f"✅ 成功生成: {file_name} | 包含 {len(batch)} 个股票")

            except Exception as e:
                print(f"❌ 生成失败: {file_name} | 错误: {str(e)}")
                if os.path.exists(file_name):
                    os.remove(file_name)

    def backtest_all_stocks_and_export(self, stock_codes: list, start_date: str, end_date: str, matype: str):
        """回测所有股票并导出信号"""
        all_signals = {}

        for stock_code in stock_codes:
            print(f"Backtesting {stock_code} from {start_date} to {end_date}")
            signals = self.backtest_strategy(stock_code, start_date, end_date)

            if signals is not None and not signals.empty:
                all_signals[stock_code] = signals
                print(f"{stock_code} Signals Count: {signals.shape[0]}")  # 添加调试输出

        if all_signals:
            self.export_signals_to_excel(all_signals, 1000, matype)
        else:
            print("⚠️ 无有效数据可导出")



    def backtest_strategy(self, symbol: str, start_date: str = "2025-01-01", end_date: str = "2025-12-31"):
        """策略回测（修复版）"""
        raw_data = self.data_manager.query_stock_data(symbol, start_date, end_date)

        if raw_data.empty:
            print(f"No data found for {symbol} between {start_date} and {end_date}")
            return None

        try:
            processed_df = raw_data.copy()
            if 'date' in processed_df.columns:
                processed_df['date'] = pd.to_datetime(processed_df['date'])
                processed_df.set_index('date', inplace=True)
            else:
                print(f"警告: {symbol} 缺少 'date' 列，无法继续处理。")
                return None

            # 生成信号
            signals = self.generate_signals(processed_df)

            # 添加必要元数据
            signals['symbol'] = symbol
            # 格式化日期並賦值（例如：轉為 'YYYY-MM-DD' 字串）
            signals['strategy_date'] = processed_df.index.strftime('%Y-%m-%d')

            if signals.empty:
                print(f"No signals generated for {symbol}.")
                return None

            print(f"Signals for {symbol}: {signals}")  # 添加调试输出
            return signals

        except Exception as e:
            print(f"回测过程中发生错误: {str(e)}")
            return None


if __name__ == "__main__":
    # 初始化数据管理器
    data_manager = StockHistoryManager()

    # 初始化交易策略
    strategy = StockTradingStrategy(data_manager)

    # 获取2025年的所有股票代码
    sz_stock_codes_2025 = data_manager.get_stock_codes_for_year(2025, ['sz'])
    # 打印获取的股票代码
    print(sz_stock_codes_2025)
    # 设置开始和结束日期
    start_date = "2025-01-01"
    end_date = "2025-12-31"
    # stock_code = "000001"
    # stock_code_list = [stock_code]
    # 回测所有股票并导出结果
    strategy.backtest_all_stocks_and_export(sz_stock_codes_2025['stock_code'], start_date, end_date, 'sz')

    # 获取2025年的所有股票代码
    sh_stock_codes_2025 = data_manager.get_stock_codes_for_year(2025, ['sh'])
    # 打印获取的股票代码
    print(sh_stock_codes_2025)
    # 设置开始和结束日期
    start_date = "2025-01-01"
    end_date = "2025-12-31"
    # stock_code = "000001"
    # stock_code_list = [stock_code]
    # 回测所有股票并导出结果
    strategy.backtest_all_stocks_and_export(sh_stock_codes_2025['stock_code'], start_date, end_date, 'sh')
